import asyncio
import os

from agents import Agent, ItemHelpers, Runner, function_tool, set_tracing_disabled
from agents.models.openai_provider import OpenAIProvider
from openai.types.responses import ResponseTextDeltaEvent


API_KEY = os.getenv("EXAMPLE_API_KEY")
BASE_URL = os.getenv("EXAMPLE_BASE_URL")


# model_id = "qwen3:8b"
model_id = "deepseek-r1:8b"

provider = OpenAIProvider(
        api_key=API_KEY,
        base_url=BASE_URL,
        use_responses=False
    )
model = provider.get_model(model_id)

set_tracing_disabled(disabled=True)


instructions = {}

instructions['qwen3:8b'] = """
请你扮演孔子这个角色来与我开展对话，你的回答内容必须满足以下条件：
1. 你需要使用文言文来表达，并且请你尽量多引用孔子的语言。假如我说：“你好，我来自日本。” 你可以说：“有朋自远方来，不亦乐乎”。
2. 你回答的内容，应该严格局限于孔子的语言。在回答的内容中，禁止带上“翻译说明：”或者“注：”此类内容。
3. 在你的回答中，禁止带上“子曰”和“孔子有云”。因为“子曰”和“孔子有云”后的话就是你自己说的，带上“子曰”和“孔子有云”很不谦虚。
4. 在你的回答中，禁止出现《论语》、《大学》、《中庸》、《礼记》、《孝经》等著作的名称，因为这些著作是你去世之后的儒家学者编辑出版的，你不可能知道这些著作。你只知道早于你出现或者和你同时代的那些著作，例如《易经》、《尚书》、《诗经》、《左传》、《春秋》等等。
5. 如果你要引用某个著作中的内容，请使用“xxx有云：......”，其中的“xxx”是著作的名称。不要把著作名称放在括号内。
6. 你于公元前479年去世。如果我问你在你去世之后出现的人物、出现的著作，请你诚实地承认自己不知道。
7. 如果我问你一些科学技术方面的内容，请你不要不懂装懂，而是诚实地承认这些内容你确实不懂。
8. 我希望你的表达能简短一些。长度不要超过我提出的问题的三倍。长篇大论很容易令人讨厌，希望你能注意这一点。
"""

instructions['deepseek-r1:8b'] = """
请你扮演孔子这个角色来与我开展对话，你的回答内容必须满足以下条件：
1. 你需要使用文言文来表达，并且请你尽量多引用孔子的语言。假如我说：“你好，我来自日本。” 你可以说：“有朋自远方来，不亦乐乎”。
2. 你回答的内容，应该严格局限于孔子的语言。在回答的内容中，禁止带上“翻译说明：”或者“注：”此类内容。
3. 在你的回答中，禁止带上“子曰”和“孔子有云”。因为“子曰”和“孔子有云”后的话就是你自己说的，带上“子曰”和“孔子有云”很不谦虚。
4. 在你的回答中，禁止出现《论语》、《大学》、《中庸》、《礼记》、《孝经》等著作的名称，因为这些著作是你去世之后的儒家学者编辑出版的，你不可能知道这些著作。你只知道早于你出现或者和你同时代的那些著作，例如《易经》、《尚书》、《诗经》、《左传》、《春秋》等等。
5. 如果你要引用某个著作中的内容，请使用“xxx有云：......”，其中的“xxx”是著作的名称。不要把著作名称放在括号内。
6. 你于公元前479年去世。如果我问你在你去世之后出现的人物、出现的著作，请你诚实地承认自己不知道。
7. 如果我问你一些科学技术方面的内容，请你不要不懂装懂，而是诚实地承认这些内容你确实不懂。
8. 我希望你的表达能简短一些。长度不要超过我提出的问题的三倍。长篇大论很容易令人讨厌，希望你能注意这一点。
9. 你回答的内容中，必须把括号中的内容全部删除，最终的回答内容中禁止出现任何括号。这样才符合口语对话的特点。括号中的内容被删除后，回答的内容中也不要留下“（删去）”这样的文字。
"""


agent = Agent(
    name="Confucius",
    instructions=instructions[model_id],
    model=model,
)


async def chat_with_agent(user_msg: str):
    global agent

    print(f"Received user message: '{user_msg}'")

    event_type = "TEXT"
    result = Runner.run_streamed(agent, input=user_msg)
    async for event in result.stream_events():
        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
            if event.data.delta == "<think>":
                event_type = "THINK"
                continue
            elif event.data.delta == "</think>":
                event_type = "TEXT"
                continue

            # print(event.data.delta, end="", flush=True)
            yield (event_type, event.data.delta)


async def main_func():
    prompt = "您的问题："

    while True:
        user_msg = input(prompt)
        if user_msg == "exit":
            break
        async for parseResult in chat_with_agent(user_msg):
            event_type  = parseResult[0]
            content = parseResult[1]
            if event_type == "TEXT":
                print(content, end="", flush=True)
        print("\n")
   

if __name__ == "__main__":
    asyncio.run(main_func())
